#include "darknet_yolo.h"

image mat_to_image(cv::Mat mat)
{
  int w = mat.cols;
  int h = mat.rows;
  int c = mat.channels();
  image im = make_image(w, h, c);

  unsigned char *data = (unsigned char *)mat.data;
  int step = mat.step;

  for (int y = 0; y < h; ++y) {
    for (int k = 0; k < c; ++k) {
      for (int x = 0; x < w; ++x) {
        //uint8_t val = mat.ptr<uint8_t>(y)[c * x + k];
        //uint8_t val = mat.at<Vec3b>(y, x).val[k];
        //im.data[k*w*h + y*w + x] = val / 255.0f;
        im.data[k*w*h + y*w + x] = data[y*step + x*c + k] / 255.0f;
      }
    }
  }
  return im;
}

DarknetYolo::~DarknetYolo()
{
  if (NULL != this->net)
  {
    free_network(*(this->net));
    this->net = NULL;
  }
}

DarknetYolo::DarknetYolo(std::string modelfile, std::string weightfile, int batchsize, int platform, float nms, float thresh)
{
  this->modelfile = modelfile;
  this->weightfile = weightfile;
  this->batchsize = batchsize;
  this->platform = platform;
  this->nms = nms;
  this->thresh = thresh;

  //load_network_custom
  char cfg[256];
  char weights[256];
  strncpy(cfg, modelfile.c_str(), 256);
  strncpy(weights, weightfile.c_str(), 256);

  network* net = load_network_custom(cfg, weights, 0, batchsize);
  this->net = net;
}

detection* DarknetYolo::Getdets(cv::Mat mat, int* nboxes)
{
  this->imwidth = mat.cols;
  this->imheight = mat.rows;

  cv::Mat sized;
  cv::resize(mat, sized, cv::Size(this->net->w,this->net->h));
  image im = mat_to_image(sized);

  float *X = im.data;
  network_predict(*(this->net), X);
  layer l = net->layers[net->n - 1];

  *nboxes = 0;
  detection *dets = get_network_boxes(this->net, this->imwidth, this->imheight, this->thresh, 0.5, 0, 1, nboxes, 0);

  if (this->nms) {
    do_nms_sort(dets, *nboxes, l.classes, this->nms);
  }

  free_image(im);
  return dets;
}

